from time import sleep
import akshare as ak
import os
import sqlite3
import pandas as pd
from datetime import datetime

conn = sqlite3.connect(os.path.join(os.getcwd(), 'index_data_cache.db'), check_same_thread=False)

# 定义一个函数来检查表是否存在
def table_exists(table_name):
    query = "SELECT name FROM sqlite_master WHERE type='table' AND name=?"
    cursor = conn.execute(query, (table_name,))
    return cursor.fetchone() is not None

def set_index_data(index_code):
    start_date: str = "19700101"
    end_date: str = datetime.now().strftime('%Y%m%d')
    df = ak.index_zh_a_hist(symbol=index_code, period="daily", start_date=start_date, end_date=end_date)
    if not df.empty:
        table_name = f"index_data_{index_code}"
        # 目标表已存在，则替换
        df.to_sql(table_name, conn, if_exists='replace', index=False)
        print(f"Data for {index_code} cached in table '{table_name}'.")

def get_index_data(index_code, start_date, end_date):
    table_name = f"index_data_{index_code}"
    is_table_exist = table_exists(table_name)
    if not is_table_exist:
        set_index_data(index_code)
        sleep(2)
    try:
        # 连接到 SQLite 数据库
        query = f"""
                SELECT *
                FROM {table_name}
                WHERE 日期 BETWEEN ? AND ?
                """
        # 使用 pandas 读取 SQL 查询结果
        df = pd.read_sql_query(query, conn, params=(start_date, end_date))
        return df
    except sqlite3.OperationalError:
        print(f"No cached data found for symbol {index_code}.")
        return None

def main():
    index_code = "000001"
    start_date = "20240606"
    end_date = "20250101"
    df = get_index_data(index_code, start_date, end_date)

if __name__ == '__main__':
    main()